#!/usr/bin/env python3
import os
import sys

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/')
)

import grpc
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections


def set_ecn_threshold(p4info_helper, switch, threshold_value):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecn_config",
        match_fields={},
        action_name="MyIngress.set_ecn_threshold",
        action_params={"threshold": threshold_value}
    )
    switch.WriteTableEntry(table_entry)
    print(f"[INFO] ECN threshold set to {threshold_value} on {switch.name}")


def install_ipv4_lpm(p4info_helper, switch, dst_ip, dst_mac, out_port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip, 32)},
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": out_port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"[{switch.name}] Rule: {dst_ip} -> {dst_mac} via port {out_port}")


def main():
    p4info_path = "build/ecn.p4.p4info.txt"
    bmv2_json_path = "build/ecn.json"

    try:
        threshold = int(input("Please input the threshold of the queue: "))
    except ValueError:
        print("[ERROR] Please enter an integer for threshold.")
        return

    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_path)

    try:
        # === 连接 s1 ===
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-ecn.txt'
        )
        s1.MasterArbitrationUpdate()
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s1, threshold)

        # === 连接 s2 ===
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-ecn.txt'
        )
        s2.MasterArbitrationUpdate()
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s2, threshold)

        # === 连接 s3 ===
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-ecn.txt'
        )
        s3.MasterArbitrationUpdate()
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s3, threshold)

        # === 安装转发规则（IP / MAC / Port）===
        # s1: h1 (1), h11 (2), s2 (3), s3 (4)
        install_ipv4_lpm(p4info_helper, s1, "10.0.1.1", "08:00:00:00:01:01", 1)
        install_ipv4_lpm(p4info_helper, s1, "10.0.1.11", "08:00:00:00:01:11", 2)
        install_ipv4_lpm(p4info_helper, s1, "10.0.2.2", "08:00:00:00:02:02", 3)  # h2 via s2
        install_ipv4_lpm(p4info_helper, s1, "10.0.2.22", "08:00:00:00:02:22", 3)
        install_ipv4_lpm(p4info_helper, s1, "10.0.3.3", "08:00:00:00:03:03", 4)  # h3 via s3

        # s2: h2 (2), h22 (1), s1 (3), s3 (4)
        install_ipv4_lpm(p4info_helper, s2, "10.0.2.2", "08:00:00:00:02:02", 2)
        install_ipv4_lpm(p4info_helper, s2, "10.0.2.22", "08:00:00:00:02:22", 1)
        install_ipv4_lpm(p4info_helper, s2, "10.0.1.1", "08:00:00:00:01:01", 3)
        install_ipv4_lpm(p4info_helper, s2, "10.0.1.11", "08:00:00:00:01:11", 3)
        install_ipv4_lpm(p4info_helper, s2, "10.0.3.3", "08:00:00:00:03:03", 4)

        # s3: h3 (1), s1 (2), s2 (3)
        install_ipv4_lpm(p4info_helper, s3, "10.0.3.3", "08:00:00:00:03:03", 1)
        install_ipv4_lpm(p4info_helper, s3, "10.0.1.1", "08:00:00:00:01:01", 2)
        install_ipv4_lpm(p4info_helper, s3, "10.0.1.11", "08:00:00:00:01:11", 2)
        install_ipv4_lpm(p4info_helper, s3, "10.0.2.2", "08:00:00:00:02:02", 3)
        install_ipv4_lpm(p4info_helper, s3, "10.0.2.22", "08:00:00:00:02:22", 3)

        print("\n[INFO] All ECN thresholds and routing rules installed.\n")

    except grpc.RpcError as e:
        printGrpcError(e)

    finally:
        ShutdownAllSwitchConnections()


if __name__ == '__main__':
    main()
